import openai
import pandas as pd
import json
import glob
from flask import Flask, request, render_template, send_file
import os
from docx import Document
import fitz
import requests
from bs4 import BeautifulSoup
import logging
import tempfile

API_KEY = "sk-ioSk7SbGiIu40lnOINybFdlyf2oOMbteB1dQ4jRwrhpqONtQ"
BASE_URL = "https://api.getaiai.com"
MODEL_NAME = "gpt-4o"

client = openai.Client(api_key=API_KEY, base_url=f"{BASE_URL}/v1")
app = Flask(__name__)

def read_files(file_path_pattern):
    file_contents = []
    for file_path in glob.glob(file_path_pattern):
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
            file_contents.append(content)
    return file_contents

def read_docx(doc_path):
    doc = Document(doc_path)
    full_text = []
    for para in doc.paragraphs:
        full_text.append(para.text)
    return '\n'.join(full_text)

def read_pdf(stream):
    document = fitz.open(stream)
    text = ""
    for page_num in range(len(document)):
        page = document.load_page(page_num)
        text += page.get_text("text")
    return text

def fetch_web_content(url):
    try:
        response = requests.get(url)
        response.raise_for_status()  # Raise an error for bad responses
        soup = BeautifulSoup(response.content, 'html.parser')
        return soup.get_text()
    except requests.RequestException as e:
        logging.error(f"Error fetching {url}: {e}")
        return ""

def check_compliance(law_contents, company_contents):
    results = []
    for company_content in company_contents:
        for law_content in law_contents:
            SYSTEM_PROMPT = f"""你将使用中文回答用户的问题，除非用户使用的是其他语言。
下面用户将提供相关法律法规片段和公司的文件，你将参考用户提供的法律法规片段和其他任何可能相关的中国法律，检查这些文件并返回不符合法规的内容
你将回复JSON格式字符串，JSON定义如下：
{{ "result" : [{{"content":"<公司文件片段>", "law":"违反《中华人民共和国XXX法》第<>条、第<>条"}}]}}
result: 检查的结果list
content: 不符合法规的内容部分，尽可能精简
law: 违反的法律法规，尽可能精简
你将输出完整的plain json字符串，不包括markdown格式等其他内容。

法律片段如下：
{law_content}

公司文件如下：
{company_content}
"""
            USER_MESSAGE = ""
            completion = client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": USER_MESSAGE},
                ],
                response_format={"type": "json_object"},
            )
            try:
                response_json = json.loads(completion.choices[0].message.content)
                results.extend(response_json.get("result", []))
            except json.JSONDecodeError:
                logging.error(f"解析结果出现错误，返回内容为：{completion.choices[0].message.content}")
    return results if results is not None else []

@app.route('/', methods=['GET'])
def index():
    return render_template('index.html')

@app.route('/check', methods=['POST'])
def check():
    if request.method == 'POST':
        law_files = request.files.getlist('law_files')
        company_files = request.files.getlist('company_files')
        company_urls = request.form.get('company_urls')

        if not law_files:
            return "请上传法律法规文件", 400

        if not (company_files or company_urls):
            return "请上传公司文件或输入公司网址", 400

        law_contents = []
        company_contents = []

        # 处理上传的法律法规文件
        for file in law_files:
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1])
            file.save(temp_file.name)
            temp_file.close()

            if file.filename.endswith('.docx'):
                law_contents.append(read_docx(temp_file.name))
            elif file.filename.endswith('.pdf'):
                with open(temp_file.name, 'rb') as pdf_file:
                    law_contents.append(read_pdf(pdf_file))
            else:
                with open(temp_file.name, 'r', encoding='utf-8') as txt_file:
                    law_contents.append(txt_file.read())

            os.remove(temp_file.name)  # 删除临时文件

        # 处理上传的公司文件
        for file in company_files:
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1])
            file.save(temp_file.name)
            temp_file.close()

            if file.filename.endswith('.docx'):
                company_contents.append(read_docx(temp_file.name))
            elif file.filename.endswith('.pdf'):
                with open(temp_file.name, 'rb') as pdf_file:
                    company_contents.append(read_pdf(pdf_file))
            else:
                with open(temp_file.name, 'r', encoding='utf-8') as txt_file:
                    company_contents.append(txt_file.read())

            os.remove(temp_file.name)

        # 处理输入的公司网址
        if company_urls:
            for url in company_urls.split(','):
                url = url.strip()
                if url:
                    web_content = fetch_web_content(url)
                    if web_content:
                        # 创建一个临时的 Document 对象来存储网页内容
                        doc = Document()
                        doc.add_paragraph(web_content)
                        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.docx')
                        doc.save(temp_file.name)
                        temp_file.close()
                        company_contents.append(read_docx(temp_file.name))
                        os.remove(temp_file.name)  # 删除临时文件

        compliance_results = check_compliance(law_contents, company_contents)

        # 将结果整理成DataFrame，准备写入Excel
        data = [{"content": item.get("content", ""), "law": item.get("law", "")} for item in compliance_results]
        df = pd.DataFrame(data)

        # 创建临时Excel文件保存结果
        excel_file = tempfile.NamedTemporaryFile(delete=False, suffix='.xlsx')
        df.to_excel(excel_file.name, index=False)
        excel_file.close()

        return render_template('result.html', results=data, temp_file=excel_file.name)

@app.route('/export-excel', methods=['GET'])
def export_excel():
    temp_file = request.args.get('temp_file')
    if not temp_file or not os.path.exists(temp_file):
        return "文件不存在", 404

    return send_file(temp_file, as_attachment=True, download_name="compliance_results.xlsx")

@app.route('/export-docx', methods=['GET'])
def export_docx():
    temp_file = request.args.get('temp_file')
    if not temp_file or not os.path.exists(temp_file):
        return "文件不存在", 404

    # 加载之前生成的数据
    df = pd.read_excel(temp_file)

    # 创建一个临时的 docx 文件保存结果
    docx_file = tempfile.NamedTemporaryFile(delete=False, suffix='.docx')
    doc = Document()
    for _, row in df.iterrows():
        doc.add_paragraph(f"Content: {row['content']}")
        doc.add_paragraph(f"Law: {row['law']}")
        doc.add_paragraph("")  # 添加空行以提高可读性

    doc.save(docx_file.name)
    docx_file.close()

    return send_file(docx_file.name, as_attachment=True, download_name="compliance_results.docx")

if __name__ == '__main__':
    app.run(debug=True)